package com.superb.common.database.utils;

import com.baomidou.mybatisplus.core.toolkit.StringPool;
import lombok.SneakyThrows;
import net.sf.jsqlparser.expression.*;
import net.sf.jsqlparser.expression.operators.arithmetic.Addition;
import net.sf.jsqlparser.expression.operators.relational.Between;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.statement.update.UpdateSet;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.*;

/**
 * 解析sql语义
 * @Author: ajie
 * @CreateTime: 2024-07-12 15:40
 */
public class SqlParserUtils {

    /**
     * @param
     * @return
     */
    @SneakyThrows
    public static Set<String> getUpdateSetItem(String sql) {
        Set<String> updateSetItemSet = new HashSet<>();
        Update update = (Update) CCJSqlParserUtil.parse(sql);
        getUpdateSetItem(update, updateSetItemSet);
        return updateSetItemSet;
    }

    /**
     * 获取update里面的set参数
     *
     * @param update
     * @param updateSetItemSet
     */
    public static void getUpdateSetItem(Update update, Set<String> updateSetItemSet) {
        List<UpdateSet> updateSets = update.getUpdateSets();
        for (UpdateSet updateSet : updateSets) {
            ExpressionList<Column> columns = updateSet.getColumns();
            for (Column column : columns) {
                if (column.getTable() == null) {
                    updateSetItemSet.add(column.getColumnName());
                } else {
                    updateSetItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
                }
            }
            for (Expression expression : columns) {
                if (expression instanceof Function) {
                    getFunction((Function) expression, updateSetItemSet);
                } else if (expression instanceof CaseExpression) {
                    getCaseExpression((CaseExpression) expression, updateSetItemSet);
                } else if (expression instanceof CastExpression) {
                    getCastExpression((CastExpression) expression, updateSetItemSet);
                } else if (expression instanceof Addition) {
                    getAddition((Addition) expression, updateSetItemSet);
                } else if (expression instanceof Column column) {
                    if (column.getTable() == null) {
                        updateSetItemSet.add(column.getColumnName());
                    } else {
                        updateSetItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
                    }
                }
            }
        }
    }


    /**
     * 获取select里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static Set<String> getSelectItem(String sql) {
        Set<String> selectItemSet = new HashSet<>();
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        getSelectItem(body, selectItemSet);
        return selectItemSet;
    }

    /**
     * 获取select里面的字段
     *
     * @param plainSelect
     * @param selectItemSet
     * @return
     */
    public static void getSelectItem(PlainSelect plainSelect, Set<String> selectItemSet) {
        List<SelectItem<?>> selectItems = plainSelect.getSelectItems();
        if (!CollectionUtils.isEmpty(selectItems)) {
            for (SelectItem selectItem : selectItems) {
                selectItem.accept(new SelectItemVisitorAdapter());
            }
        }
    }


    /**
     * 获取where里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static Set<String> getWhereItem(String sql) {
        Set<String> whereItemSet = new HashSet<>();
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        getWhereItem(body.getWhere(), whereItemSet);
        return whereItemSet;
    }

    /**
     * 获取where里面的字段
     *
     * @param
     * @return
     */
    public static void getWhereItem(Expression where, Set<String> whereItemSet) {
        if (where instanceof BinaryExpression binaryExpression) {
            Expression rightExpression = binaryExpression.getRightExpression() instanceof Parenthesis ? ((Parenthesis) binaryExpression.getRightExpression()).getExpression() : binaryExpression.getRightExpression();
            Expression leftExpression = binaryExpression.getLeftExpression() instanceof Parenthesis ? ((Parenthesis) binaryExpression.getLeftExpression()).getExpression() : binaryExpression.getLeftExpression();
            if (rightExpression instanceof Column rightColumn) {
                if (rightColumn.getTable() == null) {
                    whereItemSet.add(rightColumn.getColumnName());
                } else {
                    whereItemSet.add(rightColumn.getTable().getName() + StringPool.DOT + rightColumn.getColumnName());
                }
            } else if (rightExpression instanceof Function) {
                getFunction((Function) rightExpression, whereItemSet);
            } else if (rightExpression instanceof CaseExpression) {
                getCaseExpression((CaseExpression) rightExpression, whereItemSet);
            } else if (rightExpression instanceof CastExpression) {
                getCastExpression((CastExpression) rightExpression, whereItemSet);
            } else if (rightExpression instanceof Addition) {
                getAddition((Addition) rightExpression, whereItemSet);
            } else if (rightExpression instanceof Between) {
                getBetween((Between) rightExpression, whereItemSet);
            } else {
                getWhereItem(rightExpression, whereItemSet);
            }

            if (leftExpression instanceof Column leftColumn) {
                if (leftColumn.getTable() == null) {
                    whereItemSet.add(leftColumn.getColumnName());
                } else {
                    whereItemSet.add(leftColumn.getTable().getName() + StringPool.DOT + leftColumn.getColumnName());
                }
            } else if (leftExpression instanceof Function) {
                getFunction((Function) leftExpression, whereItemSet);
            } else if (leftExpression instanceof CaseExpression) {
                getCaseExpression((CaseExpression) leftExpression, whereItemSet);
            } else if (leftExpression instanceof CastExpression) {
                getCastExpression((CastExpression) leftExpression, whereItemSet);
            } else if (leftExpression instanceof Addition) {
                getAddition((Addition) leftExpression, whereItemSet);
            } else if (leftExpression instanceof Between) {
                getBetween((Between) leftExpression, whereItemSet);
            } else {
                getWhereItem(leftExpression, whereItemSet);
            }
        } else if (where instanceof Parenthesis) {
            getWhereItem(((Parenthesis) where).getExpression(), whereItemSet);
        } else if (where instanceof Column column) {
            if (column.getTable() == null) {
                whereItemSet.add(column.getColumnName());
            } else {
                whereItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
            }
        } else if (where instanceof Function) {
            getFunction((Function) where, whereItemSet);
        } else if (where instanceof CaseExpression) {
            getCaseExpression((CaseExpression) where, whereItemSet);
        } else if (where instanceof CastExpression) {
            getCastExpression((CastExpression) where, whereItemSet);
        }
    }

    /**
     * 获取orderby里面的参数
     *
     * @param sql
     * @return
     */
    @SneakyThrows
    public static Set<String> getOrderByItem(String sql) {
        Set<String> orderByItemSet = new HashSet<>();
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        getOrderByItem(body, orderByItemSet);
        return orderByItemSet;
    }

    /**
     * 获取orderby里面的参数
     *
     * @param plainSelect
     * @param orderByItemSet
     */
    public static void getOrderByItem(PlainSelect plainSelect, Set<String> orderByItemSet) {
        List<OrderByElement> orderByElements = plainSelect.getOrderByElements();
        if (!CollectionUtils.isEmpty(orderByElements)) {
            for (OrderByElement orderByElement : orderByElements) {
                orderByElement.accept(new OrderByVisitorAdapter() {
                    @Override
                    public void visit(OrderByElement item) {
                        if (item.getExpression() instanceof Function) {
                            getFunction((Function) item.getExpression(), orderByItemSet);
                        } else if (item.getExpression() instanceof CaseExpression) {
                            getCaseExpression((CaseExpression) item.getExpression(), orderByItemSet);
                        } else if (item.getExpression() instanceof CastExpression) {
                            getCastExpression((CastExpression) item.getExpression(), orderByItemSet);
                        } else if (item.getExpression() instanceof Addition) {
                            getAddition((Addition) item.getExpression(), orderByItemSet);
                        } else if (item.getExpression() instanceof Column) {
                            Column column = (Column) item.getExpression();
                            if (column.getTable() == null) {
                                orderByItemSet.add(column.getColumnName());
                            } else {
                                orderByItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
                            }
                        }
                    }
                });
            }
        }
    }

    /**
     * 获取groupBy里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static Set<String> getGroupBy(String sql) {
        Set<String> groupByItemSet = new HashSet<>();
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        getGroupBy(body, groupByItemSet);
        return groupByItemSet;
    }

    /**
     * 获取groupBy里面的字段
     *
     * @param
     * @return
     */
    public static void getGroupBy(PlainSelect plainSelect, Set<String> groupByItemSet) {
        GroupByElement groupByElement = plainSelect.getGroupBy();
        if (groupByElement != null) {
            List<Expression> expressions = groupByElement.getGroupByExpressions();
            for (Expression item : expressions) {
                if (item instanceof Column) {
                    Column column = (Column) item;
                    if (column.getTable() == null) {
                        groupByItemSet.add(column.getColumnName());
                    } else {
                        groupByItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
                    }
                }
            }
        }
    }

    /**
     * 设置替换update set里面的字段
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static String setUpdateSetItem(String sql, Map<String, String> map) {
        Update update = (Update) CCJSqlParserUtil.parse(sql);
        setUpdateSetItem(update, map);
        return update.toString();
    }

    /**
     * 设置替换update set里面的字段
     *
     * @param update
     * @param map
     * @return
     */
    public static void setUpdateSetItem(Update update, Map<String, String> map) {
        List<UpdateSet> updateSets = update.getUpdateSets();
        for (UpdateSet updateSet : updateSets) {
            List<Column> columns = updateSet.getColumns();
            for (Column column : columns) {
                column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
            }
            ExpressionList<?> expressions = updateSet.getValues();
            for (Expression expression : expressions) {
                if (expression instanceof Function) {
                    setFunction((Function) expression, map);
                } else if (expression instanceof CaseExpression) {
                    setCaseExpression((CaseExpression) expression, map);
                } else if (expression instanceof CastExpression) {
                    setCastExpression((CastExpression) expression, map);
                } else if (expression instanceof Addition) {
                    setAddition((Addition) expression, map);
                } else if (expression instanceof Column column) {
                    column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
                }
            }
        }
    }

    /**
     * 设置select里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static String setSelectItem(String sql, Map<String, String> map) {
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = selectStatement.getPlainSelect();
        setSelectItem(body, map);
        return body.toString();
    }

    /**
     * 设置替换select 里面的字段
     *
     * @param plainSelect
     * @param map
     * @return
     */
    public static void setSelectItem(PlainSelect plainSelect, Map<String, String> map) {
        List<SelectItem<?>> selectItems = plainSelect.getSelectItems();
        if (!CollectionUtils.isEmpty(selectItems)) {
            for (SelectItem selectItem : selectItems) {
                selectItem.accept(new SelectItemVisitorAdapter());
            }
        }
    }

    /**
     * 获取where里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static String setWhere(String sql, Map<String, String> map) {
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        setWhereItem(body.getWhere(), map);
        return body.toString();
    }

    /**
     * 设置替换where里面的字段
     *
     * @param
     * @return
     */
    public static void setWhereItem(Expression where, Map<String, String> map) {
        if (where instanceof BinaryExpression) {
            BinaryExpression binaryExpression = (BinaryExpression) where;
            Expression rightExpression = binaryExpression.getRightExpression() instanceof Parenthesis ? ((Parenthesis) binaryExpression.getRightExpression()).getExpression() : binaryExpression.getRightExpression();
            Expression leftExpression = binaryExpression.getLeftExpression() instanceof Parenthesis ? ((Parenthesis) binaryExpression.getLeftExpression()).getExpression() : binaryExpression.getLeftExpression();
            if (rightExpression instanceof Column) {
                Column rightColumn = (Column) rightExpression;
                rightColumn.setColumnName(Optional.ofNullable(map.get(rightColumn.getColumnName())).orElse(rightColumn.getColumnName()));
            } else if (rightExpression instanceof Function) {
                setFunction((Function) rightExpression, map);
            } else if (rightExpression instanceof CaseExpression) {
                setCaseExpression((CaseExpression) rightExpression, map);
            } else if (rightExpression instanceof CastExpression) {
                setCastExpression((CastExpression) rightExpression, map);
            } else if (rightExpression instanceof Addition) {
                setAddition((Addition) rightExpression, map);
            } else if (rightExpression instanceof Between) {
                setBetween((Between) rightExpression, map);
            } else {
                setWhereItem(rightExpression, map);
            }
            if (leftExpression instanceof Column) {
                Column leftColumn = (Column) leftExpression;
                leftColumn.setColumnName(Optional.ofNullable(map.get(leftColumn.getColumnName())).orElse(leftColumn.getColumnName()));
            } else if (leftExpression instanceof Function) {
                setFunction((Function) leftExpression, map);
            } else if (leftExpression instanceof CaseExpression) {
                setCaseExpression((CaseExpression) leftExpression, map);
            } else if (leftExpression instanceof CastExpression) {
                setCastExpression((CastExpression) leftExpression, map);
            } else if (leftExpression instanceof Addition) {
                setAddition((Addition) leftExpression, map);
            } else if (leftExpression instanceof Between) {
                setBetween((Between) leftExpression, map);
            } else {
                setWhereItem(leftExpression, map);
            }
        } else if (where instanceof Parenthesis) {
            setWhereItem(((Parenthesis) where).getExpression(), map);
        } else if (where instanceof Column) {
            Column column = (Column) where;
            column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
        } else if (where instanceof Function) {
            setFunction((Function) where, map);
        } else if (where instanceof CaseExpression) {
            setCaseExpression((CaseExpression) where, map);
        } else if (where instanceof CastExpression) {
            setCastExpression((CastExpression) where, map);
        } else if (where instanceof Addition) {
            setAddition((Addition) where, map);
        }
    }


    /**
     * 设置orderBy里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static String setOrderByItem(String sql, Map<String, String> map) {
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        setOrderByItem(body, map);
        return body.toString();
    }

    /**
     * 设置替换orderBy里面的字段
     *
     * @param plainSelect
     * @param map
     * @return
     */
    public static void setOrderByItem(PlainSelect plainSelect, Map<String, String> map) {
        List<OrderByElement> orderByElements = plainSelect.getOrderByElements();
        if (!CollectionUtils.isEmpty(orderByElements)) {
            for (OrderByElement orderByElement : orderByElements) {
                orderByElement.accept(new OrderByVisitorAdapter() {
                    @Override
                    public void visit(OrderByElement item) {
                        if (item.getExpression() instanceof Function) {
                            setFunction((Function) item.getExpression(), map);
                        } else if (item.getExpression() instanceof CaseExpression) {
                            setCaseExpression((CaseExpression) item.getExpression(), map);
                        } else if (item.getExpression() instanceof CastExpression) {
                            setCastExpression((CastExpression) item.getExpression(), map);
                        } else if (item.getExpression() instanceof Addition) {
                            setAddition((Addition) item.getExpression(), map);
                        } else if (item.getExpression() instanceof Column) {
                            Column column = (Column) item.getExpression();
                            column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
                        }
                    }
                });
            }
        }
    }

    /**
     * 设置groupBy里面的参数
     *
     * @param
     * @return
     */
    @SneakyThrows
    public static String setGroupBy(String sql, Map<String, String> map) {
        Select selectStatement = (Select) CCJSqlParserUtil.parse(sql);
        PlainSelect body = (PlainSelect) selectStatement.getSelectBody();
        setGroupBy(body, map);
        return body.toString();
    }

    /**
     * 设置groupBy里面的字段
     *
     * @param
     * @return
     */
    public static void setGroupBy(PlainSelect plainSelect, Map<String, String> map) {
        GroupByElement groupByElement = plainSelect.getGroupBy();
        if (groupByElement != null) {
            ExpressionList<Expression> expressions = groupByElement.getGroupByExpressionList();
            for (Expression item : expressions) {
                if (item instanceof Column column) {
                    column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
                }
            }
        }
    }

    /**
     * 获取function里面的字段
     *
     * @param function
     * @param selectItemSet
     * @return
     */
    public static void getFunction(Function function, Set<String> selectItemSet) {
        ExpressionList<?> parameters = function.getParameters();
        if (parameters == null || parameters.getASTNode() == null) {
            return;
        }
        parameters.forEach(data -> {
            if (data instanceof Function) {
                getFunction((Function) data, selectItemSet);
            } else if (data instanceof Column column) {
                if (column.getTable() == null) {
                    selectItemSet.add(column.getColumnName());
                } else {
                    selectItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
                }
            } else {
                getWhereItem(data, selectItemSet);
            }
        });

    }

    /**
     * 获取CaseExpression里面的字段
     *
     * @param caseExpression
     * @param selectItemSet
     * @return
     */
    public static void getCaseExpression(CaseExpression caseExpression, Set<String> selectItemSet) {
        if (caseExpression.getWhenClauses() == null) {
            return;
        }
        List<WhenClause> list = caseExpression.getWhenClauses();
        list.forEach(data -> {
            if (data != null) {
                getWhereItem(data.getWhenExpression(), selectItemSet);
            }
        });
    }

    /**
     * 获取CastExpression里的字段
     *
     * @param castExpression
     * @param selectItemSet
     */
    public static void getCastExpression(CastExpression castExpression, Set<String> selectItemSet) {
        if (castExpression.getLeftExpression() == null) {
            return;
        }
        if (castExpression.getLeftExpression() instanceof Column column) {
            if (column.getTable() == null) {
                selectItemSet.add(column.getColumnName());
            } else {
                selectItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
            }
        }
    }

    /**
     * 获取Addition里的字段
     *
     * @param addition
     * @param selectItemSet
     */
    public static void getAddition(Addition addition, Set<String> selectItemSet) {
        if (addition.getLeftExpression() == null) {
            return;
        }
        if (addition.getLeftExpression() instanceof Function) {
            getFunction((Function) addition.getLeftExpression(), selectItemSet);
        }
        if (addition.getLeftExpression() instanceof Column column) {
            if (column.getTable() == null) {
                selectItemSet.add(column.getColumnName());
            } else {
                selectItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
            }
        }
    }

    /**
     * 获取Between里的字段
     *
     * @param between
     * @param selectItemSet
     */
    public static void getBetween(Between between, Set<String> selectItemSet) {
        if (between.getLeftExpression() == null) {
            return;
        }
        if (between.getLeftExpression() instanceof Function) {
            getFunction((Function) between.getLeftExpression(), selectItemSet);
        }
        if (between.getLeftExpression() instanceof Column column) {
            if (column.getTable() == null) {
                selectItemSet.add(column.getColumnName());
            } else {
                selectItemSet.add(column.getTable().getName() + StringPool.DOT + column.getColumnName());
            }
        }
    }

    /**
     * 设置替换 function里面的字段
     *
     * @param function
     * @param map
     * @return
     */
    public static void setFunction(Function function, Map<String, String> map) {
        ExpressionList<?> parameters = function.getParameters();
        if (parameters == null || parameters.getASTNode() == null) {
            return;
        }
        parameters.forEach(data -> {
            if (data instanceof Function) {
                setFunction((Function) data, map);
            } else if (data instanceof Column column) {
                column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
            } else {
                setWhereItem(data, map);
            }
        });

    }

    /**
     * 设置替换 CaseExpression里面的字段
     *
     * @param caseExpression
     * @param map
     * @return
     */
    public static void setCaseExpression(CaseExpression caseExpression, Map<String, String> map) {
        if (caseExpression.getWhenClauses() == null) {
            return;
        }
        List<WhenClause> list = caseExpression.getWhenClauses();
        list.forEach(data -> {
            if (data != null) {
                setWhereItem(data.getWhenExpression(), map);
            }
        });
    }

    /**
     * 设置替换 CastExpression里面的字段
     *
     * @param castExpression
     * @param map
     */
    public static void setCastExpression(CastExpression castExpression, Map<String, String> map) {
        if (castExpression.getLeftExpression() == null) {
            return;
        }
        if (castExpression.getLeftExpression() instanceof Column rightColumn) {
            rightColumn.setColumnName(Optional.ofNullable(map.get(rightColumn.getColumnName())).orElse(rightColumn.getColumnName()));
        }
    }

    /**
     * 设置Addition里的字段
     *
     * @param addition
     * @param map
     */
    public static void setAddition(Addition addition, Map<String, String> map) {
        if (addition.getLeftExpression() == null) {
            return;
        }
        if (addition.getLeftExpression() instanceof Function) {
            setFunction((Function) addition.getLeftExpression(), map);
        }
        if (addition.getLeftExpression() instanceof Column column) {
            column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
        }
    }

    /**
     * 获取Between里的字段
     *
     * @param between
     * @param map
     */
    public static void setBetween(Between between, Map<String, String> map) {
        if (between.getLeftExpression() == null) {
            return;
        }
        if (between.getLeftExpression() instanceof Function) {
            setFunction((Function) between.getLeftExpression(), map);
        }
        if (between.getLeftExpression() instanceof Column column) {
            column.setColumnName(Optional.ofNullable(map.get(column.getColumnName())).orElse(column.getColumnName()));
        }
    }

}
